package structure.bbst.avl;




import java.util.Comparator;

/**
 * 项目名称: vital
 * 包 名 称: com.bbst.avl
 * 类 名 称: AvlTree
 * 类 描 述: TODO
 * 创建时间: 2020/6/7 11:58 上午
 * 创 建 人: Justice
 * 搜索、添加、删除的时间复杂度是 O(logn)
 */
public class AvlTree<E> extends BinarySearchTree<E> {
    private static class AvlNode<E> extends Node<E> {
        // 叶子结点默认高度为1
        int height = 1;

        public AvlNode(E element, Node<E> parent) {
            super(element, parent);
        }

        // 获取结点平衡因子，结点左右子树的高度差
        public int balanceFactor() {
            int leftHeight = left == null ? 0 : ((AvlNode<E>) left).height;
            int rightHeight = right == null ? 0 : ((AvlNode<E>) right).height;
            return leftHeight - rightHeight;
        }

        // 添加新结点之后，更新结点高度
        public void updateHeight() {
            int leftHeight = left == null ? 0 : ((AvlNode<E>) left).height;
            int rightHeight = right == null ? 0 : ((AvlNode<E>) right).height;
            height = 1 + Math.max(leftHeight, rightHeight);
        }

        // 返回高子结点
        public Node<E> tallerChild() {
            int leftHeight = left == null ? 0 : ((AvlNode<E>) left).height;
            int rightHeight = right == null ? 0 : ((AvlNode<E>) right).height;
            if (leftHeight > rightHeight) return left;
            if (leftHeight < rightHeight) return right;
            return isLeftChild() ? left : right;
        }

        @Override
        public String toString() {
            String parentString = "null";
            if (parent != null) {
                parentString = parent.element.toString();
            }
            return element + "_p(" + parentString + ")_h(" + height + ")";

        }
    }

    // 判断结点是否平衡
    private boolean isBalanced(Node<E> node) {
        return Math.abs(((AvlNode<E>) node).balanceFactor()) <= 1;
    }

    // 更新结点高度
    private void updateHeight(Node<E> node) {
        ((AvlNode<E>) node).updateHeight();
    }

    public AvlTree() {
        this(null);
    }

    public AvlTree(Comparator<E> comparator) {
        super(comparator);
    }

    @Override
    protected void afterAdd(Node<E> node) {
        // 添加结点之后需要调整
        while ((node = node.parent) != null) {
            if (isBalanced(node)) {
                // 更新结点高度
                updateHeight(node);
            } else {
                // 发现不平衡结点进行调整
                reBalance(node);
                // 整颗树恢复平衡之后
                break;
            }
        }
    }

    // 恢复平衡 高度最低的不平衡结点
    private void reBalance2(Node<E> grandparent) {
        Node<E> parent = ((AvlNode<E>) grandparent).tallerChild();
        Node<E> node = ((AvlNode<E>) parent).tallerChild();
        if (parent.isLeftChild()) {
            if (node.isLeftChild()) {
                // LL
                rotateRight(grandparent);
            } else {
                // LR
                rotateLeft(parent);
                rotateRight(grandparent);
            }
        } else {
            if (node.isRightChild()) {
                // RR
                rotateLeft(grandparent);
            } else {
                // RL
                rotateRight(parent);
                rotateLeft(grandparent);
            }
        }
    }

    // 恢复平衡 高度最低的不平衡结点
    private void reBalance(Node<E> grandparent) {
        Node<E> parent = ((AvlNode<E>) grandparent).tallerChild();
        Node<E> node = ((AvlNode<E>) parent).tallerChild();
        if (parent.isLeftChild()) { // L
            if (node.isLeftChild()) { // LL
                rotate(grandparent, node, node.right, parent, parent.right, grandparent);
            } else { // LR
                rotate(grandparent, parent, node.left, node, node.right, grandparent);
            }
        } else { // R
            if (node.isLeftChild()) { // RL
                rotate(grandparent, grandparent, node.left, node, node.right, parent);
            } else { // RR
                rotate(grandparent, grandparent, parent.left, parent, node.left, node);
            }
        }
    }

    private void rotate(Node<E> r, Node<E> b, Node<E> c, Node<E> d, Node<E> e, Node<E> f) {
        // 让d成为子树根结点
        d.parent = r.parent;
        if (r.isLeftChild()) {
            r.parent.left = d;
        } else if (r.isRightChild()) {
            r.parent.right = d;
        } else {
            root = d;
        }

        b.right = c;
        if (c != null) {
            c.parent = b;
        }
        updateHeight(b);

        f.left = e;
        if (e != null) {
            e.parent = f;
        }
        updateHeight(f);

        d.left = b;
        d.right = f;
        b.parent = d;
        f.parent = d;
        updateHeight(d);
    }

    // 左旋转
    private void rotateLeft(Node<E> grandparent) {
        Node<E> parent = grandparent.right;
        Node<E> child = parent.left;
        grandparent.right = child;
        parent.left = grandparent;
        afterRotate(grandparent, parent, child);
    }

    // 右旋转
    private void rotateRight(Node<E> grandparent) {
        Node<E> parent = grandparent.left;
        Node<E> child = parent.right;
        grandparent.left = child;
        parent.right = grandparent;
        afterRotate(grandparent, parent, child);
    }

    // 旋转之后调整关系
    private void afterRotate(Node<E> grandparent, Node<E> parent, Node<E> child) {
        // 首先让parent成为子树的根结点
        parent.parent = grandparent.parent;
        // grandparent左右子树指向parent
        if (grandparent.isLeftChild()) {
            grandparent.parent.left = parent;
        } else if (grandparent.isRightChild()) {
            grandparent.parent.right = parent;
        } else {
            // grandparent本来是就是根结点
            root = parent;
        }
        // 更新child的parent
        if (child != null) {
            child.parent = grandparent;
        }
        // 更新grandparent的parent
        grandparent.parent = parent;
        // 更新高度
        updateHeight(grandparent);
        updateHeight(parent);
    }

    @Override
    protected void afterRemove(Node<E> node) {
        while ((node = node.parent) != null) {
            if (isBalanced(node)) {
                // 更新高度
                updateHeight(node);
            } else {
                // 恢复平衡
                reBalance(node);
            }
        }
    }

    @Override
    protected Node<E> createNode(E element, Node<E> parent) {
        // 创建avl树的结点
        return new AvlNode<E>(element, parent);
    }
}
